{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TI Sitara x AWS NEO Image Classification Example\n",
    "\n",
    "1. [Introduction](#Introduction)\n",
    "2. [Compile model using NEO](#Compile-model-using-NEO)\n",
    "3. [Inference on device](#Inference-on-device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "This notebook will demo how to compile pretrained TensorFlow Mobilenet V2 model using AWS Neo for TI Sitara AM57xx. At last, we will deploy compiled model to device and do inference using the Neo Deep Learning Runtime.\n",
    "\n",
    "To get started, we need to set up the environment for AWS S3 permissions, configurations, and so on. Please refer to [Configuration](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use Boto 3, first import it and tell it what service we are going to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "sm = boto3.client('sagemaker', region_name='us-west-2')\n",
    "s3 = boto3.client('s3')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "account = boto3.client('sts').get_caller_identity().get('Account')\n",
    "print(account)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get the pretrained tensorflow mobilenet V2 model \n",
    "\n",
    "The pretrained tensorflow models can be found [here](https://github.com/tensorflow/models/tree/master/research/slim)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://storage.googleapis.com/mobilenet_v2/checkpoints/mobilenet_v2_1.0_224.tgz "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = 'mobilenet_v2_1.0_224'\n",
    "model = model_name + '.tar.gz'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -xf mobilenet_v2_1.0_224.tgz -C .\n",
    "!tar -czf $model mobilenet_v2_1.0_224_frozen.pb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Upload model to S3 bucket"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a S3 bucket `ti-demo` to store pretrained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket = 'sitara-demo'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if boto3.resource('s3').Bucket(bucket) not in boto3.resource('s3').buckets.all():\n",
    "    s3.create_bucket(\n",
    "        Bucket=bucket,\n",
    "        CreateBucketConfiguration={\n",
    "            'LocationConstraint': 'us-west-2'\n",
    "        }\n",
    "    )\n",
    "else:\n",
    "    print('Bucket %s already exists' %bucket)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3.upload_file(model, bucket, model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create IAM role\n",
    "In order to use the sagemaker service and have access to S3 bucket, we need to create a IAM role."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam = boto3.client('iam')\n",
    "role_name = 'TI-demo-test-role'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = {\n",
    "    'Statement': [{\n",
    "        'Action': 'sts:AssumeRole',\n",
    "        'Effect': 'Allow',\n",
    "        'Principal': {'Service': 'sagemaker.amazonaws.com'}},\n",
    "        ],  \n",
    "     'Version': '2012-10-17'}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "roles = iam.list_roles()\n",
    "role_arn = None\n",
    "for role in roles['Roles']:\n",
    "    if role['RoleName'] == role_name:\n",
    "        role_arn = role['Arn']\n",
    "        \n",
    "if role_arn == None:\n",
    "    new_role = iam.create_role(\n",
    "        AssumeRolePolicyDocument=json.dumps(policy),\n",
    "        Path='/',\n",
    "        RoleName=role_name,\n",
    "    )\n",
    "    role_arn = new_role['Role']['Arn']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam.attach_role_policy(\n",
    "    RoleName=role_name,\n",
    "    PolicyArn='arn:aws:iam::aws:policy/AmazonSageMakerFullAccess'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iam.attach_role_policy(\n",
    "    RoleName=role_name,\n",
    "    PolicyArn='arn:aws:iam::aws:policy/AmazonS3FullAccess'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile model using NEO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_output_location = 's3://{}/output'.format(bucket)\n",
    "data_shape = '{\"input\":[1,224,224,3]}' # you can change batch_size as you need, here use 1\n",
    "framework = 'tensorflow'\n",
    "target_device = 'sitara_am57x'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "compilation_job_name = 'ti-demo'+ str(time.time()).split('.')[0]\n",
    "print('Compilation job for %s started' % compilation_job_name)\n",
    "\n",
    "response = sm.create_compilation_job(\n",
    "        CompilationJobName=compilation_job_name,\n",
    "        RoleArn=role_arn,\n",
    "        InputConfig={\n",
    "            'S3Uri': 's3://{}/{}'.format(bucket, model),\n",
    "            'DataInputConfig': data_shape,\n",
    "            'Framework': framework.upper()\n",
    "        },\n",
    "        OutputConfig={\n",
    "            'S3OutputLocation': s3_output_location,\n",
    "            'TargetDevice': target_device \n",
    "        },\n",
    "        StoppingCondition={\n",
    "            'MaxRuntimeInSeconds': 900\n",
    "        }\n",
    "    )\n",
    "\n",
    "print(response)\n",
    "\n",
    "# Poll every 30 sec\n",
    "while True:\n",
    "    response = sm.describe_compilation_job(CompilationJobName=compilation_job_name)\n",
    "    if response['CompilationJobStatus'] == 'COMPLETED':\n",
    "        break\n",
    "    elif response['CompilationJobStatus'] == 'FAILED':\n",
    "        raise RuntimeError('Compilation failed')\n",
    "    print('Compiling ...')\n",
    "    time.sleep(30)\n",
    "print('Done!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference on device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download compiled model from S3 to device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_path = 'output/{}-{}.tar.gz'.format(model_name, target_device)\n",
    "neo_compiled_model = 'compiled-'+ model\n",
    "s3.download_file(bucket, object_path, neo_compiled_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir compiled_model\n",
    "!tar -xf $neo_compiled_model -C ./compiled_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp ./compiled_model/*.bin .\n",
    "!cp ./compiled_model/*.cfg ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use DLR to read compiled model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dlr import DLRModel\n",
    "import numpy as np\n",
    "import time\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model\n",
    "model_path = \"./compiled_model\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cpu'\n",
    "model = DLRModel(model_path, device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download an image to prepare for predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!wget -O test.jpg http://www.vision.caltech.edu/Image_Datasets/Caltech256/images/080.frog/080_0001.jpg\n",
    "\n",
    "import json\n",
    "import numpy as np\n",
    "    \n",
    "file_name = \"test.jpg\"\n",
    "\n",
    "# test image\n",
    "from IPython.display import Image\n",
    "Image(file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Image pre-process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import PIL.Image\n",
    "image = PIL.Image.open(file_name)\n",
    "\n",
    "batch_size = 1\n",
    "\n",
    "# Resize\n",
    "image = np.asarray(image.resize((224, 224)))\n",
    "\n",
    "# Normalize\n",
    "image = image/128 - 1\n",
    "\n",
    "image = np.concatenate([image[np.newaxis, :, :]]*batch_size)\n",
    "\n",
    "print(image.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference and prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#flatten within a input array\n",
    "input_data = {'input': image}\n",
    "\n",
    "# dry run\n",
    "for _ in range(10):\n",
    "    model.run(input_data)\n",
    "\n",
    "print('Testing inference on mobilenet_v2...')\n",
    "start_time = time.time()\n",
    "out = model.run(input_data) #need to be a list of input arrays matching input names\n",
    "index = np.argmax(out[0][0,:])\n",
    "prob = np.amax(out[0][0,:])\n",
    "print('inference time is ' + str((time.time()-start_time)/batch_size) + ' seconds')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load names for ImageNet classes\n",
    "object_categories = {}\n",
    "with open(\"imagenet1000_clsidx_to_labels.txt\", \"r\") as f:\n",
    "    for line in f:\n",
    "        key, val = line.strip().split(':')\n",
    "        object_categories[key] = val\n",
    "print(index)\n",
    "print(\"Result: label - \" + object_categories[str(index)] + \" probability - \" + str(prob))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
